#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
将PyTorch模型转换为safetensors格式
"""
import os
import sys
import torch
from safetensors.torch import save_file

def convert_model(model_path):
    """将PyTorch模型转换为safetensors格式"""
    try:
        print(f"正在转换模型: {model_path}")
        
        # 检查模型文件是否存在
        pytorch_model_path = os.path.join(model_path, "pytorch_model.bin")
        if not os.path.exists(pytorch_model_path):
            print(f"错误: 未找到 {pytorch_model_path}")
            return False
            
        # 加载PyTorch模型
        print("正在加载PyTorch模型...")
        state_dict = torch.load(pytorch_model_path, map_location="cpu")
        
        # 确保所有张量都是连续的
        print("正在处理张量连续性...")
        for key in state_dict:
            if isinstance(state_dict[key], torch.Tensor) and not state_dict[key].is_contiguous():
                state_dict[key] = state_dict[key].contiguous()
        
        # 保存为safetensors格式
        safetensors_path = os.path.join(model_path, "model.safetensors")
        print(f"正在保存为safetensors格式: {safetensors_path}")
        save_file(state_dict, safetensors_path)
        
        print("模型转换完成!")
        return True
        
    except Exception as e:
        print(f"转换失败: {str(e)}")
        import traceback
        traceback.print_exc()
        return False

def main():
    """主函数"""
    if len(sys.argv) < 2:
        print("用法: python convert_model_to_safetensors.py <model_path>")
        sys.exit(1)
        
    model_path = sys.argv[1]
    
    if not os.path.exists(model_path):
        print(f"错误: 模型路径不存在: {model_path}")
        sys.exit(1)
        
    success = convert_model(model_path)
    sys.exit(0 if success else 1)

if __name__ == "__main__":
    main()
